# -*- coding: utf-8 -*-

# -------------------------------
# Author    : linlingwei
# Email     : 951035650@qq.com
# File Name : hairDyn.py
# Date      : 2020/3/13 
# IDE       : PyCharm
# Version   : 1.0.1
# -------------------------------

"""
import sys
path = r"D:\Code\x6_scripts\code" #替换成脚本所在路径
if path not in sys.path:
    sys.path.append(path)
import hairDyn as hairDyn
reload(hairDyn)
hairDyn.makeIkHair("joint1","joint13")
"""

import maya.cmds as mc
import maya.mel as mel

def makeIkHair(starJoint,endJoint):
    
    # todo create Curve on joints
    curveName = createCurveOnChain(starJoint,endJoint,"test")
    
    # todo create Hair system duplicateCurve follicle
    hairSystemShape,follicle,curveShape = makeIkHair_makeCurvesDynamicHairs([curveName])
    FollicleShape = mc.listRelatives(follicle,s=True)
    
    # todo create SplineHandle for selected chain and the duplicated curve
    ik = mc.ikHandle(sol="ikSplineSolver",ccv=False,pcv=False,snc=True,sj=starJoint,ee=endJoint,c=curveShape)[0]

    # todo create duplicate the chain for animation
    duplSkel = makeIkHair_duplicateSkeletonAndConnect("Test_",[starJoint])

    # todo create SkinCluster between duplicateCurve and animation chain
    skinCluster = mc.skinCluster(duplSkel,curveName)

    # todo
    #  add and connect attributes on root annimation
    #  Chain controling stiffness damping and gravity
    #  iteration startFrame collisions

    controller = mel.eval('$controller = `curve -d 1 -p -2 2 2 -p 2 2 2 -p 2 -2 2 -p -2 -2 2 -p -2 2 2 -p -2 2 -2 -p -2 -2 -2 -p -2 -2 2 -p -2 -2 -2 -p 2 -2 -2 -p 2 2 -2 -p -2 2 -2 -p 2 2 -2 -p 2 2 2 -p 2 -2 2 -p 2 -2 -2 -k 0 -k 1 -k 2 -k 3 -k 4 -k 5 -k 6 -k 7 -k 8 -k 9 -k 10 -k 11 -k 12 -k 13 -k 14 -k 15 `;')
    shapes = mc.listRelatives(controller,s=True)
    shape = mc.rename(shapes[0],duplSkel[0]+"Shape")
    mc.parent(shape,duplSkel[0],s=True,r=True)
    mc.delete(controller)
    
    mc.addAttr(duplSkel[0],ln="___DYNAMICS___",at="float")
    mc.setAttr(duplSkel[0]+".___DYNAMICS___",e=True,k=True,l=True)
    mc.addAttr(duplSkel[0],ln="chainAttach",at="enum",en="No Attach:Base:Tip:Both End",dv=1)
    mc.setAttr(duplSkel[0]+".chainAttach",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainStartEnveloppe",at="float",min=0,max=1,dv=1)
    mc.setAttr(duplSkel[0]+".chainStartEnveloppe",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainStartFrame",at="float",dv=1)
    mc.setAttr(duplSkel[0]+".chainStartFrame",e=True,k=True)
    
    mc.addAttr(duplSkel[0],ln="___BEHAVIOR___",at="float")
    mc.setAttr(duplSkel[0]+".___BEHAVIOR___",e=True,k=True,l=True)
    mc.addAttr(duplSkel[0],ln="chainStiffness",at="float",dv=1,min=0,max=1)
    mc.setAttr(duplSkel[0]+".chainStiffness",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainDamping",at="float",dv=.2,min=0)
    mc.setAttr(duplSkel[0]+".chainDamping",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainGravity",at="float",dv=1,min=0)
    mc.setAttr(duplSkel[0]+".chainGravity",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainIteration",at="long",dv=1,min=0)
    mc.setAttr(duplSkel[0]+".chainIteration",e=True,k=True)

    mc.addAttr(duplSkel[0],ln="___COLLISIONS___",at="float")
    mc.setAttr(duplSkel[0]+".___COLLISIONS___",e=True,k=True,l=True)
    mc.addAttr(duplSkel[0],ln="chainCollide",at="bool",dv=0)
    mc.setAttr(duplSkel[0]+".chainCollide",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainWidthBase",at="float",dv=1,min=0)
    mc.setAttr(duplSkel[0]+".chainWidthBase",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainWidthExtremity",at="float",dv=1,min=0)
    mc.setAttr(duplSkel[0]+".chainWidthExtremity",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainCollideGround",at="bool",dv=0)
    mc.setAttr(duplSkel[0]+".chainCollideGround",e=True,k=True)
    mc.addAttr(duplSkel[0],ln="chainCollideGroundHeight",at="float",dv=0)
    mc.setAttr(duplSkel[0]+".chainCollideGroundHeight",e=True,k=True)
    
    # ---
    mc.connectAttr(duplSkel[0]+".chainCollideGroundHeight",ik + ".ikBlend",f=True)
    mc.connectAttr(duplSkel[0]+".chainAttach",FollicleShape[0] + ".pointLock",f=True)
    mc.connectAttr(duplSkel[0]+".chainStartFrame",hairSystemShape + ".startFrame",f=True)
    mc.connectAttr(duplSkel[0]+".chainStiffness",hairSystemShape + ".stiffness",f=True)
    mc.connectAttr(duplSkel[0]+".chainDamping",hairSystemShape + ".damp",f=True)
    mc.connectAttr(duplSkel[0]+".chainGravity",hairSystemShape + ".gravity",f=True)
    mc.connectAttr(duplSkel[0]+".chainIteration",hairSystemShape + ".iterations",f=True)
    mc.connectAttr(duplSkel[0]+".chainCollide",hairSystemShape + ".collide",f=True)
    mc.connectAttr(duplSkel[0]+".chainWidthBase",hairSystemShape + ".clumpWidth",f=True)
    mc.connectAttr(duplSkel[0]+".chainWidthExtremity",hairSystemShape + ".clumpWidthScale[1].clumpWidthScale_FloatValue",f=True)
    mc.connectAttr(duplSkel[0]+".chainCollideGround",hairSystemShape + ".collideGround",f=True)
    mc.connectAttr(duplSkel[0]+".chainCollideGroundHeight",hairSystemShape + ".groundHeight",f=True)
    
    loc = mc.spaceLocator()[0]
    track_RealTime = mc.rename(loc,"TRACK_RealTime_" + loc[0])
    mc.pointConstraint(duplSkel[-2],track_RealTime,)
    connections = mc.listConnections(hairSystemShape+".currentTime",p=True,c=True)
    mc.disconnectAttr(connections[1],connections[0])
    mc.expression(s="if (frame != {create}.startFrame)\r\n"
                    "{create}.currentTime={create}.currentTime+1"
                    "+{trt}.tx-{trt}.tx +{trt}.ty-{trt}.ty +{trt}.tz-{trt}.tz"
                    "+{dup}.chainWidthBase-{dup}.chainWidthBase +{dup}.chainWidthExtremity-{dup}.chainWidthExtremity +{dup}.chainGravity-{dup}.chainGravity;\n"\
                    "else\n\t{create}.currentTime = {create}.startFrame;".format(create=hairSystemShape,trt=track_RealTime,dup=duplSkel[0]))


    # todo
    #  set default presets for dynHairChain
    
    mc.setAttr(hairSystemShape+".drawCollideWidth",1)
    mc.setAttr(hairSystemShape+".widthDrawSkip",0)
    mc.setAttr(FollicleShape[0]+".degree",1)
    parent = mc.listRelatives(starJoint,p=True)
    if parent:
        mc.parent(duplSkel[0],parent)
    if not mc.objExists("dynHairChain"):
        mc.group(em=True,name="dynHairChain")
    if not mc.objExists(starJoint+"_HairControls"):
        grp = mc.group(hairSystemShape,follicle,curveShape,ik,track_RealTime,name = starJoint+"_HairControls",)
    else:
        mc.parent(hairSystemShape,follicle,curveShape,ik,starJoint+"_HairControls")
    mc.parent(grp,"dynHairChain")

    return

def createCurveOnChain(starJoint,endJoint,crvName):
    listJoints = find_jointChain(endJoint, starJoint)
    listJoints.reverse()
    positions = [mc.xform(lj,q=True,ws=True,t=True) for lj in listJoints]
    curveName = mc.rename(mc.curve(p=positions,d=1),crvName)
    return curveName+"Shape"

def find_jointChain(child, father):
    joint = child
    array = [joint]
    i = 0
    while joint != father:
        nextJoint = mc.listRelatives(joint,p=True,f=True,typ="joint")
        if nextJoint:
            array.append(nextJoint[0])
            joint = nextJoint
        if i >10000: break
        i+=1
    return array

def makeIkHair_makeCurvesDynamicHairs(curves):
    created = []
    hsystems = mel.eval("getSelectedHairSystems;")
    hsys = ""
    if len(hsystems)>0:
        hsys = hsystems[0]

    madeHairCurve = False
    hsysGroup = ""
    hsysOutputHairGroup = ""
    
    if not curves: return
    
    for crv in curves:
        con = mc.listConnections(crv+".worldSpace[0]",sh=True)
        attachedToHairCurve = False
        if con:
            for c in con :
                if mc.nodeType(c) == "follicle":
                    attachedToHairCurve = True
                    continue
        if attachedToHairCurve : continue
        intermediateObject = mc.getAttr(crv+".io")
        if intermediateObject != 0: continue

        madeHairCurve = True
        if hsys == "":
            hsys = mc.createNode("hairSystem")
            created.append(hsys)
            mc.removeMultiInstance(hsys+".stiffnessScale[1]",b=True)
            mc.setAttr(hsys+".clumpWidth",0)
            mc.connectAttr("time1.outTime",hsys+".currentTime")
        hsysParent = mel.eval("string $hsysParent[] = `listTransforms {}`;".format(hsys))
        if len(hsysParent) > 0 :
            if hsysGroup == "" :
                hsysGroup = hsysParent[0]+"Follicles"
                if not mc.objExists(hsysGroup):
                    hsysGroup = mc.group(em=True,name = hsysGroup)
            if hsysOutputHairGroup == "":
                hsysOutputHairGroup = hsysParent[0]+"OutputCurves"
                if not mc.objExists(hsysOutputHairGroup):
                    hsysOutputHairGroup = mc.group(em=True,name = hsysOutputHairGroup)

        COS = False
        relatives = mc.listRelatives(crv,ap=True,p=True)
        for testMe in relatives:
            if mc.nodeType(testMe) == "transform":
                testList = mc.listRelatives(testMe,ap=True,p=True)
                if testList:
                    for subTest in testList:
                        if mc.nodeType(subTest) == "curveVarGroup":
                            subTestList = mc.listRelatives(subTest,ap=True,p=True)
                            for stl in subTestList:
                                if mc.nodeType(stl) == "nurbsSurface":
                                    COS = True
                                    parent=stl
                                    break
                        else:
                            if mc.nodeType(subTest) == "nurbsSurface":
                                COS = True
                                parent=subTest
                                break
            else:
                if mc.nodeType(testMe) == "nurbsSurface":
                    COS = True
                    parent=testMe
                    break
                    
        if COS :
            cFS = mc.createNode("curveFromSurfaceCoS")
            created.append(cFS)
            mc.connectAttr(crv+".ws[0]",cFS+".curveOnSurface")
            mc.connectAttr(parent+".ws[0]",cFS+".inputSurface")
            crv = cFS
            
        surf = ""
        doRest = False
        tmp = makeIkHair_createHairCurveNode(hsys,surf,0.0, 0.0, 0, True, True,
                                             doRest, False, crv, 1.0,[0],
                                             hsysGroup,hsysOutputHairGroup)
        hname = tmp[0]
        created.append(hname)
        created.append(tmp[1])
        if len(hname) > 0:
            mc.setAttr(hname+".pointLock",3)
            mc.setAttr(hname+".restPose",1)
    if mc.objExists(hsysGroup):
        shapes = mc.ls(hsysGroup,s=True,dag=True)
        if len(shapes) == 0:
            mc.delete(hsysGroup)
    if mc.objExists(hsysOutputHairGroup):
        shapes = mc.ls(hsysOutputHairGroup,s=True,dag=True)
        if len(shapes) == 0:
            mc.delete(hsysOutputHairGroup)
    if madeHairCurve:
        mc.select(hsys)
    else:
        pass
    
    #hsys
    return created


def makeIkHair_createHairCurveNode(hsys,surface,u,v,numCvs,
                                   doOut,doStart,doRest,
                                   isPassive,startCurve,length,
                                   endHairSystemIndex,
                                   hsysGroup,hsysOutHairGroup):
    
    doStartCurve = False
    COS = False
    if not hsys: return
    hair = mc.createNode("follicle")
    mc.setAttr(hair+".parameterU",u)
    mc.setAttr(hair+".parameterV",v)
    tforms = mel.eval("string $tforms[] = `listTransforms {}`;".format(hair))
    hairDag = tforms[0]
    attachedToSurface = False
    if surface != "" and mc.objExists(surface):
        nType = mc.nodeType(surface)
        mc.connectAttr(surface+".worldMatrix[0]",hair+".inputWorldMatrix")
        if nType == "nurbsSurface":
            mc.connectAttr(surface+".local",hair+".inputSurface")
        elif nType == "mesh":
            mc.connectAttr(surface+".outMesh",hair+".inputMesh")
            currentUVSet = mc.polyUVSet(surface,q=True,currentUVSet=True)
            mc.setAttr(hair+".mapSetName",currentUVSet[0],type="string",)
            isValidUv = mc.getAttr(hair+".validUv")
            if not isValidUv:
                mc.delete(hair)
                return
    
        mc.connectAttr(hair+".outTranslate",hairDag+".t")
        mc.connectAttr(hair+".outRotate",hairDag+".r")
        mc.setAttr(hairDag+".t",l=True)
        mc.setAttr(hairDag+".r",l=True)
    else:
        mc.setAttr(hair+".startDirection",1)
    if doStart:
        if startCurve != "":
            if mc.objExists(startCurve):
                typ = mc.nodeType(startCurve)
                if typ == "nurbsCurve":
                    doStartCurve = True
                elif typ == "curveFromSurfaceCoS":
                    doStartCurve = True
                    COS = True
                if doStartCurve:
                    curveCvs = mc.getAttr(startCurve+".cp")
                    if curveCvs < 3:
                        mc.setAttr( hair+".degree",1)
                else: pass
        if not doStartCurve:
            cmd = polyLineCmd(numCvs,length)
            if cmd == "":
                mc.delete(hair)
                return
            if numCvs < 3:
                mc.setAttr(hair+".degree",1)
    elif doRest:
        cmd = polyLineCmd(numCvs,length)
        if cmd == "":
            mc.delete(hair)
            return
    if isPassive and not doStart and numCvs < 3:
        mc.setAttr(hair+".degree",1)
    
    if hsysGroup != "":
        mc.parent(hairDag,hsysGroup,r=True)
    
    if doStart:
        if doStartCurve:
            if COS:
                curve = startCurve
            else:
                tforms = mel.eval("$tforms = `listTransforms {}`;".format(startCurve))
                curve = tforms[0]
                tforms = mel.eval("$tforms = `listTransforms {}`;".format(curve))
                if len(tforms) > 0:
                    curveParent = tforms[0]
                    mc.parent( hairDag,curveParent,r=True)
        else:
            curve = polyLineCmd(numCvs,length)
            if isPassive:
                mel.eval('initHairCurveDisplay( {}, "passive" )'.format(curve))
            else:
                mel.eval('initHairCurveDisplay( {}, "start" )'.format(curve))
        if COS :
            mc.connectAttr(curve+".oc",hair+".startPosition")
        else:
            mc.parent(curve,hairDag,r=True)
            mc.connectAttr(curve+".worldSpace[0]",hair+".startPosition")
    if doRest:
        rest = ""
        if doStartCurve:
            rest = mc.duplicate(curve,rr=True)
        else:
            rest = mel.eval()#todo
        mel.eval('initHairCurveDisplay( {}, "rest" );'.format(rest))
        mc.parent(rest,hairDag,r=True)
        mc.connectAttr(rest+".ws[0]",hair+".restPosition")
        mc.setAttr(hair+".restPose",3)
    if isPassive:
        mc.setAttr(hair+".simulationMethod",1)

    endHairSystemIndex = getNextFreeMultiIndex(hsys+".inputHair",endHairSystemIndex[0])
    hairIndex = endHairSystemIndex
    endHairSystemIndex += 1
    mc.connectAttr(hair+".outHair",hsys+".inputHair[{}]".format(hairIndex))
    if doOut:
        mc.connectAttr(hsys+".outputHair[{}]".format(hairIndex),hair+".currentPosition")
        crv = mc.createNode("nurbsCurve")
        mc.connectAttr(hair+".outCurve",crv+".create")
        print crv
        if isPassive:
            mel.eval('initHairCurveDisplay( "{}", "passive" );'.format(crv))
        else:
            mel.eval('initHairCurveDisplay( "{}", "current" );'.format(crv))
        if hsysOutHairGroup != "":
            tforms = mel.eval('$tforms = `listTransforms {}`;'.format(crv))
            mc.parent(tforms[0],hsysOutHairGroup,r=True)
    return hairDag,crv
def getNextFreeMultiIndex(attr,start):
    for i in range(start,10000000):
        con = mc.connectionInfo(attr+"[{}]".format(i),sfd=True)
        if not con:
            return i
    return 0
def polyLineCmd(numCvs,curveLength):
    if numCvs < 2: return
    
    fac = curveLength/(numCvs*-1.)
    pts = []
    for i in range(numCvs):
        pts.append([0,0,i*fac])
    crv = mc.curve(p=pts,d=1)
    return crv

def makeIkHair_duplicateSkeletonAndConnect(prefix,joints):
    
    for startJoint1 in joints:
        nodes = mc.duplicate(startJoint1,rr=True,rc=True)
        startJoint2 = nodes[0]
        joints1 = mc.listRelatives(startJoint1,ad=True)
        joints2 = mc.listRelatives(startJoint2,ad=True)
        
        for i,j1 in enumerate(joints1):
            if mc.objectType(joints2[i],isType="joint"):
                mc.rename(joints2[i],prefix+j1)
                mc.connectAttr(prefix+j1+".r",j1+".r")
            else:
                mc.delete(joints2[i])
        mc.connectAttr(startJoint2+".t",startJoint1+".t")
        mc.connectAttr(startJoint2+".r",startJoint1+".r")
        mc.connectAttr(startJoint2+".s",startJoint1+".s")
        
        startJoint2 = mc.rename(startJoint2,prefix+startJoint1)
        mel.eval('catch (`parent -w {}`);'.format(startJoint2))
        mc.select(startJoint2,r=True)
        mc.select(hi=True)
        relist = mc.ls(sl=True,l=True)
        return relist

    return
if __name__ == '__main__':
    pass